{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb425bf5",
   "metadata": {},
   "source": [
    "### Video Description with Qwen3-Omni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b92d046f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['VLLM_USE_V1'] = '0'\n",
    "os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'\n",
    "os.environ[\"VLLM_LOGGING_LEVEL\"] = \"ERROR\"\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"0\"\n",
    "import torch\n",
    "import warnings\n",
    "import numpy as np\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
    "warnings.filterwarnings('ignore', category=FutureWarning)\n",
    "warnings.filterwarnings('ignore', category=UserWarning)\n",
    "\n",
    "from qwen_omni_utils import process_mm_info\n",
    "from transformers import Qwen3OmniMoeProcessor\n",
    "\n",
    "def _load_model_processor():\n",
    "    if USE_TRANSFORMERS:\n",
    "        from transformers import Qwen3OmniMoeForConditionalGeneration\n",
    "        if TRANSFORMERS_USE_FLASH_ATTN2:\n",
    "            model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(MODEL_PATH,\n",
    "                                                                         dtype='auto',\n",
    "                                                                         attn_implementation='flash_attention_2',\n",
    "                                                                         device_map=\"auto\")\n",
    "        else:\n",
    "            model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(MODEL_PATH, device_map=\"auto\", dtype='auto')\n",
    "    else:\n",
    "        from vllm import LLM\n",
    "        model = LLM(\n",
    "            model=MODEL_PATH, trust_remote_code=True, gpu_memory_utilization=0.95,\n",
    "            tensor_parallel_size=torch.cuda.device_count(),\n",
    "            limit_mm_per_prompt={'image': 1, 'video': 3, 'audio': 3},\n",
    "            max_num_seqs=1,\n",
    "            max_model_len=32768,\n",
    "            seed=1234,\n",
    "        )\n",
    "\n",
    "    processor = Qwen3OmniMoeProcessor.from_pretrained(MODEL_PATH)\n",
    "    return model, processor\n",
    "\n",
    "def run_model(model, processor, messages, return_audio, use_audio_in_video):\n",
    "    if USE_TRANSFORMERS:\n",
    "        text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)\n",
    "        audios, images, videos = process_mm_info(messages, use_audio_in_video=use_audio_in_video)\n",
    "        inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors=\"pt\", padding=True, use_audio_in_video=use_audio_in_video)\n",
    "        inputs = inputs.to(model.device).to(model.dtype)\n",
    "        text_ids, audio = model.generate(**inputs, \n",
    "                                            thinker_return_dict_in_generate=True,\n",
    "                                            thinker_max_new_tokens=8192, \n",
    "                                            thinker_do_sample=False,\n",
    "                                            speaker=\"Ethan\", \n",
    "                                            use_audio_in_video=use_audio_in_video,\n",
    "                                            return_audio=return_audio)\n",
    "        response = processor.batch_decode(text_ids.sequences[:, inputs[\"input_ids\"].shape[1] :], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]\n",
    "        if audio is not None:\n",
    "            audio = np.array(audio.reshape(-1).detach().cpu().numpy() * 32767).astype(np.int16)\n",
    "        return response, audio\n",
    "    else:\n",
    "        from vllm import SamplingParams\n",
    "        sampling_params = SamplingParams(temperature=1e-2, top_p=0.1, top_k=1, max_tokens=8192)\n",
    "        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "        audios, images, videos = process_mm_info(messages, use_audio_in_video=use_audio_in_video)\n",
    "        inputs = {'prompt': text, 'multi_modal_data': {}, \"mm_processor_kwargs\": {\"use_audio_in_video\": use_audio_in_video}}\n",
    "        if images is not None: inputs['multi_modal_data']['image'] = images\n",
    "        if videos is not None: inputs['multi_modal_data']['video'] = videos\n",
    "        if audios is not None: inputs['multi_modal_data']['audio'] = audios\n",
    "        outputs = model.generate(inputs, sampling_params=sampling_params)\n",
    "        response = outputs[0].outputs[0].text\n",
    "        return response, None\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d37dcedc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section', 'interleaved'}\n",
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section', 'interleaved'}\n",
      "`torch_dtype` is deprecated! Use `dtype` instead!\n",
      "2025-09-14 20:56:01,784\tINFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n",
      "You are attempting to use Flash Attention 2 without specifying a torch dtype. This might lead to unexpected behaviour\n",
      "Loading safetensors checkpoint shards:   0% Completed | 0/15 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards:   7% Completed | 1/15 [00:00<00:09,  1.50it/s]\n",
      "Loading safetensors checkpoint shards:  20% Completed | 3/15 [00:01<00:06,  1.77it/s]\n",
      "Loading safetensors checkpoint shards:  27% Completed | 4/15 [00:02<00:08,  1.27it/s]\n",
      "Loading safetensors checkpoint shards:  33% Completed | 5/15 [00:03<00:05,  1.76it/s]\n",
      "Loading safetensors checkpoint shards:  40% Completed | 6/15 [00:04<00:06,  1.30it/s]\n",
      "Loading safetensors checkpoint shards:  47% Completed | 7/15 [00:05<00:07,  1.10it/s]\n",
      "Loading safetensors checkpoint shards:  53% Completed | 8/15 [00:06<00:06,  1.02it/s]\n",
      "Loading safetensors checkpoint shards:  60% Completed | 9/15 [00:07<00:06,  1.03s/it]\n",
      "Loading safetensors checkpoint shards:  67% Completed | 10/15 [00:08<00:05,  1.09s/it]\n",
      "Loading safetensors checkpoint shards:  73% Completed | 11/15 [00:10<00:04,  1.13s/it]\n",
      "Loading safetensors checkpoint shards:  80% Completed | 12/15 [00:11<00:03,  1.14s/it]\n",
      "Loading safetensors checkpoint shards:  87% Completed | 13/15 [00:12<00:02,  1.17s/it]\n",
      "Loading safetensors checkpoint shards:  93% Completed | 14/15 [00:13<00:01,  1.18s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 15/15 [00:14<00:00,  1.18s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 15/15 [00:14<00:00,  1.00it/s]\n",
      "\n",
      "The image processor of type `Qwen2VLImageProcessor` is now loaded as a fast processor by default, even if the model checkpoint was saved with a slow processor. This is a breaking change and may produce slightly different outputs. To continue using the slow processor, instantiate this class with `use_fast=False`. Note that this behavior will be extended to all models in a future release.\n",
      "Unused or unrecognized kwargs: images.\n",
      "Capturing CUDA graph shapes: 100%|██████████| 1/1 [00:00<00:00,  1.67it/s]\n"
     ]
    }
   ],
   "source": [
    "import librosa\n",
    "import audioread\n",
    "\n",
    "from IPython.display import Video\n",
    "from IPython.display import Image\n",
    "from IPython.display import Audio\n",
    "\n",
    "MODEL_PATH = \"Qwen/Qwen3-Omni-30B-A3B-Instruct\"\n",
    "# MODEL_PATH = \"Qwen/Qwen3-Omni-30B-A3B-Thinking\"\n",
    "\n",
    "USE_TRANSFORMERS = False\n",
    "TRANSFORMERS_USE_FLASH_ATTN2 = True\n",
    "\n",
    "model, processor = _load_model_processor()\n",
    "\n",
    "USE_AUDIO_IN_VIDEO = False\n",
    "RETURN_AUDIO = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0b1330dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video src=\"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/cookbook/video1.mp4\" controls  width=\"640\"  height=\"360\">\n",
       "      Your browser does not support the <code>video</code> element.\n",
       "    </video>"
      ],
      "text/plain": [
       "<IPython.core.display.Video object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "qwen-vl-utils using torchvision to read video.\n",
      "Adding requests:   0%|          | 0/1 [00:00<?, ?it/s]Unused or unrecognized kwargs: images.\n",
      "Adding requests: 100%|██████████| 1/1 [00:06<00:00,  6.19s/it]\n",
      "Processed prompts: 100%|██████████| 1/1 [00:06<00:00,  6.59s/it, est. speed input: 1342.97 toks/s, output: 15.63 toks/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The video showcases a series of impressive acrobatic performances by various gymnasts, highlighting their strength, balance, and coordination. The performances take place in a large indoor arena with a blue mat, surrounded by spectators and judges. The gymnasts, dressed in colorful and elaborate costumes, execute a variety of complex maneuvers, including lifts, balances, and flips. The routines are synchronized and fluid, demonstrating the athletes' skill and artistry. The audience watches intently, adding to the atmosphere of excitement and anticipation.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "video_path = \"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/cookbook/video1.mp4\"\n",
    "\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"video\", \"video\": video_path},\n",
    "            {\"type\": \"text\", \"text\": \"Describe the video.\"} \n",
    "        ]\n",
    "    }\n",
    "]\n",
    "\n",
    "display(Video(video_path, width=640, height=360))\n",
    "\n",
    "response, audio = run_model(model=model, messages=messages, processor=processor, return_audio=RETURN_AUDIO, use_audio_in_video=USE_AUDIO_IN_VIDEO)\n",
    "\n",
    "print(response)\n",
    "if audio is not None:\n",
    "    display(Audio(audio, rate=24000))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qwen3_omni_public",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
